# -*- coding: utf-8 -*-
"""
------------------------------------------------------------------------------
    File Name:  lenet_demo
    Author   :  wanwei1029
    Date     :  2019/1/9
    Desc     :
------------------------------------------------------------------------------
"""
from keras.models import Sequential
from keras.layers.convolutional import Conv2D, MaxPooling2D
from keras.layers.core import Activation, Dense, Flatten
from keras.utils import np_utils
from keras.optimizers import Adam
from keras import backend as K
import numpy as np
import matplotlib.pyplot as plt
import samp.keras_learning.dl_action.chapter01.minist_demo as md

NB_EPOCH = 20
BATCH_SIZE = 128
VERBOSE = 2
NB_CLASSES = 10
OPTIMIZER = Adam()
VALIDATION_SPLIT = 0.2
RESHAPED = 784
IMG_ROWS, IMG_COLS = 28, 28
INPUT_SHAPE = (1, IMG_ROWS, IMG_COLS)


class LeNet:
    @staticmethod
    def build(input_shape, classes):
        """
        以inut_shape为28*28为例，经过第一层卷积，变化28*28*20.
        池化后，变化14*14*20，再经过一层卷积，变化14*14*50.
        池化后，变成7*7*50，扁平化后，就是2450，然后全连接再输出。
        :param input_shape:
        :param classes:
        :return:
        """
        model = Sequential()
        """
        Conv2D 默认步长为1，如果padding为same，卷积后大小为：原长度/步长，默认步长为1则大小不变，
        没指定则默认为valid，不填充卷积后的大小为：（原大小-卷积核+1）/步长。
        """
        model.add(Conv2D(20, 5, padding="same", input_shape=input_shape))
        model.add(Activation("relu"))
        # 池化，如果步长为空的话，会默认为pool_size，一般最好能整除最好，像这种2，2的，就是直接大小缩小一半。
        model.add(MaxPooling2D(pool_size=(2, 2), strides=(2, 2)))

        model.add(Conv2D(50, 5, padding="same"))
        model.add(Activation("relu"))
        model.add(MaxPooling2D(pool_size=(2, 2), strides=(2, 2)))

        model.add(Flatten())
        # 扁平化以便进行全连接。
        model.add(Dense(500))
        model.add(Activation("relu"))

        model.add(Dense(classes))
        model.add(Activation("softmax"))
        return model


def lenet_demo():
    (x_train, y_train), (x_test, y_test) = md.load_mnist_data()
    print(x_train.shape)
    print(y_train.shape)
    K.set_image_dim_ordering("th")
    x_train = x_train.astype('float32')
    x_test = x_test.astype('float32')
    x_train /= 255
    x_test /= 255

    x_train = x_train[:, np.newaxis, :, :]
    x_test = x_test[:, np.newaxis, :, :]
    print(x_train.shape)

    y_train = np_utils.to_categorical(y_train, NB_CLASSES)
    y_test = np_utils.to_categorical(y_test, NB_CLASSES)

    model = LeNet.build(INPUT_SHAPE, NB_CLASSES)
    model.summary()
    model.compile(loss='categorical_crossentropy', optimizer=OPTIMIZER, metrics=['accuracy'])

    history = model.fit(x_train, y_train, batch_size=BATCH_SIZE, epochs=NB_EPOCH, verbose=VERBOSE,
                        validation_split=VALIDATION_SPLIT)

    score = model.evaluate(x_test, y_test, verbose=VERBOSE)

    print("Test score:", score[0])
    print("Test accuracy:", score[1])

    # 打印训练过程
    print(history.history.keys())
    # 打印准确率数据
    plt.plot(history.history['acc'])
    plt.plot(history.history['val_acc'])
    plt.title("model accuracy")
    plt.ylabel('accuracy')
    plt.xlabel('epoch')
    plt.legend(['train', 'test'], loc='upper left')
    plt.show()

    # 打印损失函数历史数据
    plt.plot(history.history['loss'])
    plt.plot(history.history['val_loss'])
    plt.title('model loss')
    plt.ylabel('loss')
    plt.xlabel('epoch')
    plt.legend(['train', 'test'], loc='upper left')
    plt.show()

def demo():
    """
    """
    lenet_demo()


if __name__ == '__main__':
    test_method = "demo"
    if test_method == "demo":
        demo()
